from FSaugClass import FSAugTrain, FSAugVal
from LMaugClass import LMAugTrain, LMAugVal
from LMaugClass import drop_shitdata
from sklearn.model_selection import train_test_split
from FSdataset import idx2attr_map
import pandas as pd
import os
from FSdataset_test import DataProvider


def prepare_FSLM_dataProvider(FSroot_path, LMroot_path,
                              FSselect_AttrIdx = range(8)
                              ):
    # prepare FS dataset
    all_pd = pd.read_csv(os.path.join(FSroot_path, 'base/Annotations/label.csv'),
                         header=None, names=['ImageName', 'AttrKey', 'AttrValues'])
    all_pd['ImageName'] = all_pd['ImageName'].apply(lambda x: os.path.join('base', x))

    train_pd, val_pd = train_test_split(all_pd, test_size=0.1, random_state=37,
                                        stratify=all_pd['AttrKey'])

    # select part
    select_AttrKey = [idx2attr_map[x] for x in FSselect_AttrIdx]
    train_pd = train_pd[train_pd['AttrKey'].apply(lambda x: True if x in select_AttrKey else False)]
    val_pd = val_pd[val_pd['AttrKey'].apply(lambda x: True if x in select_AttrKey else False)]

    add_skirt = pd.read_csv(os.path.join(FSroot_path, 'base/Annotations/data_add_skirt_legth.csv'),
                            header=None, names=['ImageName', 'AttrKey', 'AttrValues'])
    add_skirt['ImageName'] = add_skirt['ImageName'].apply(lambda x: os.path.join('web', x))
    train_pd = pd.concat([train_pd, add_skirt], axis=0, ignore_index=True)

    FStrain_pd = train_pd
    FSval_pd = val_pd

    print FStrain_pd.shape, FSval_pd.shape


    #  prepare LM dataset
    annotation = pd.read_csv(os.path.join(LMroot_path, "train/Annotations/train.csv"))
    annotation['image_id'] = annotation['image_id'].apply(lambda x: os.path.join(LMroot_path,'train', x))
    warmup = pd.read_csv(os.path.join(LMroot_path, "warm/Annotations/annotations.csv"))
    warmup['image_id'] = warmup['image_id'].apply(lambda x: os.path.join(LMroot_path, 'warm', x))


    train_pd, val_pd = train_test_split(annotation, test_size=0.1, random_state=42,
                                        stratify=annotation['image_category'])

    val_pd = drop_shitdata(val_pd.copy())
    train_pd = pd.concat([train_pd, warmup],axis=0,ignore_index=True)
    train_pd.index = range(train_pd.shape[0])

    LMtrain_pd = train_pd
    LMval_pd = val_pd

    print LMtrain_pd.shape, LMval_pd.shape


    FSanno_pd = {'train': FStrain_pd, 'val':FSval_pd}
    LManno_pd = {'train': LMtrain_pd, 'val': LMval_pd}
    FSAug = {'train': FSAugTrain, 'val': FSAugVal}
    LMAug = {'train': LMAugTrain, 'val': LMAugVal}
    FSbatch_size = {'train': 8, 'val': 2}
    LMbatch_size = {'train': 8, 'val': 2}

    data_provider = DataProvider(
        FSroot_path, FSanno_pd, FSAug, FSselect_AttrIdx,
        LManno_pd, LMAug,
        FSbatch_size, LMbatch_size
    )
    return data_provider